package com.gel.springai.redis.starter;

import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.data.redis.core.RedisTemplate;

import java.util.List;
import java.util.stream.Collectors;

public class RedisChatMemory implements ChatMemory {

    private final RedisTemplate<String, SerializableMessage> redisTemplate;

    public RedisChatMemory(RedisTemplate<String, SerializableMessage> redisTemplate) {
        this.redisTemplate = redisTemplate;
    }

    @Override
    public void add(String conversationId, Message message) {
        redisTemplate.opsForList().leftPush(conversationId, new SerializableMessage(message));
    }

    @Override
    public void add(String conversationId, List<Message> messages) {
        List<SerializableMessage> serializableMessages=messages.stream()
                .map(SerializableMessage::new)
                .collect(Collectors.toList());
        redisTemplate.opsForList().leftPushAll(conversationId, serializableMessages);
    }

    @Override
    public List<Message> get(String conversationId, int lastN) {
        List<SerializableMessage> serializableMessages = redisTemplate.opsForList().range(conversationId, 0, lastN - 1);
        if(serializableMessages==null||serializableMessages.isEmpty()){
            return List.of();
        }
        return serializableMessages.stream()
                .map(SerializableMessage::toMessage)
                .collect(Collectors.toList());
    }

    @Override
    public void clear(String conversationId) {
        redisTemplate.delete(conversationId);
    }
}
